# 分布对抗模块
import torch


def DAM(XA, XB, alpha=0.01):
    # batch, classes = XA.shape
    # for i in range(batch):
    #     for j in range(classes):
    #         XAi = XA[i][j]
    #         XBi = XB[i][j]
    #         XA[i][j] = XAi if XAi > XBi else (2 * XAi - XBi - alpha)
    XA = 2 * XA - XB - alpha
    return XA


if __name__ == '__main__':
    xa = torch.ones(2,751)
    xb = torch.zeros(2,751)
    print(xa)
    xa = DAM(xa,xb)
    print(xa)